package com.rhb.rhbaiagent.chatmemory;

import cn.hutool.core.collection.CollectionUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import com.rhb.rhbaiagent.model.dao.ConversationMemoryDAO;
import com.rhb.rhbaiagent.model.domain.ConversationMemory;
import com.rhb.rhbaiagent.model.enmu.MessageTypeEnum;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.model.Media;
import org.springframework.core.io.ClassPathResource;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.stereotype.Component;
import org.springframework.util.MimeTypeUtils;
import org.yaml.snakeyaml.nodes.CollectionNode;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.alibaba.cloud.ai.dashscope.api.DashScopeApi.ChatCompletionMessage.Role.TOOL;

/**
 * 基于MySQL数据库的对话记忆
 */
@Component
public class MySqlChatMemory implements ChatMemory {
    
    // 全局变量用于存储当前处理的图片路径信息
    private static final ThreadLocal<String> currentImagePath = new ThreadLocal<>();
    
    /**
     * 设置当前线程的图片路径
     * @param imagePath 图片路径
     */
    public static void setCurrentImagePath(String imagePath) {
        currentImagePath.set(imagePath);
    }
    
    /**
     * 获取当前线程的图片路径
     * @return 图片路径
     */
    public static String getCurrentImagePath() {
        return currentImagePath.get();
    }
    
    /**
     * 清除当前线程的图片路径
     */
    public static void clearCurrentImagePath() {
        currentImagePath.remove();
    }

    private final ConversationMemoryDAO conversationMemoryDAO;

    public MySqlChatMemory(ConversationMemoryDAO conversationMemoryDAO) {
        this.conversationMemoryDAO = conversationMemoryDAO;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        Gson gson = new Gson();
        List<ConversationMemory> conversationMemories = messages.stream().map(message -> {
            String messageType = message.getMessageType().getValue();
            String mes = "";
            StringBuilder pathInfo = new StringBuilder();
            // 处理包含图片的消息，只保存路径而不是完整的图片数据
            if (message instanceof UserMessage userMessage && userMessage.getMedia() != null && !userMessage.getMedia().isEmpty()) {
                // 创建简化的消息对象（不包含完整的媒体数据）
                ((UserMessage) message).getMedia().clear();
                UserMessage simplifiedMessage = new UserMessage(userMessage.getText());
                mes = gson.toJson(simplifiedMessage);

                // 从ThreadLocal获取图片路径信息
                String imagePath = getCurrentImagePath();
                if (imagePath != null) {
                    // 添加路径信息到序列化消息
                    mes = mes + "MEDIA_PATH:" + imagePath + ";";
                    // 清除ThreadLocal中的路径信息
                    clearCurrentImagePath();
                }
            } else {
                // 普通消息直接序列化
                mes = gson.toJson(message);
            }

            return ConversationMemory.builder().conversationId(conversationId)
                    .Type(messageType)
                    .memory(mes)
                    .build();
        }).toList();

        conversationMemoryDAO.saveBatch(conversationMemories);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<ConversationMemory> messages = conversationMemoryDAO.getMessages(conversationId);
        if(CollectionUtil.isEmpty(messages)){
            return List.of();
        }
        return messages.stream()
                .skip(Math.max(0, messages.size() - lastN))
                .map(this::getMessage)
                .collect(Collectors.toList());
    }

    private Message getMessage(ConversationMemory conversationMemory){
        String memory = conversationMemory.getMemory();
        Gson gson = new Gson();

        // 检查是否包含媒体路径信息
        if (memory.contains("MEDIA_PATH:")) {
            // 分离JSON消息和路径信息
            String[] parts = memory.split("MEDIA_PATH:");
            String jsonPart = parts[0];

            // 反序列化基础消息
            Message baseMessage = (Message) gson.fromJson(jsonPart, MessageTypeEnum.fromValue(conversationMemory.getType()).getClazz());

            // 如果是UserMessage，恢复媒体内容
            if (baseMessage instanceof UserMessage userMessage) {
                List<Media> mediaList = new ArrayList<>();

                // 解析所有媒体路径
                for (int i = 1; i < parts.length; i++) {
                    String pathPart = parts[i];
                    String path = pathPart.split(";")[0]; // 移除分号

                    // 创建Media对象
                    Media media = new Media(MimeTypeUtils.IMAGE_PNG, new ClassPathResource(path));
                    mediaList.add(media);
                }

                // 创建包含媒体的UserMessage
                return new UserMessage(userMessage.getText(), mediaList);
            }

            return baseMessage;
        } else {
            // 普通消息直接反序列化
            return (Message) gson.fromJson(memory, MessageTypeEnum.fromValue(conversationMemory.getType()).getClazz());
        }
    }


    @Override
    public void clear(String conversationId) {
        conversationMemoryDAO.deleteMemory(conversationId);
    }


}